package org.example.cases.sf1;

import com.antgroup.geaflow.api.graph.compute.VertexCentricCompute;
import com.antgroup.geaflow.api.graph.function.vc.VertexCentricCombineFunction;
import com.antgroup.geaflow.api.graph.function.vc.VertexCentricComputeFunction;
import com.antgroup.geaflow.model.graph.edge.IEdge;
import org.example.function.AbstractVcFunc;

import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;

public class Case2Algorithm {

    public static class Case2VCAlgorithm extends VertexCentricCompute<Long, Double, Double, Long> {

        public Case2VCAlgorithm(long iterations) {
            super(iterations);
        }

        @Override
        public VertexCentricComputeFunction<Long, Double, Double, Long> getComputeFunction() {
            return new Case2VertexCentricComputeFunction();
        }

        @Override
        public VertexCentricCombineFunction<Long> getCombineFunction() {
            return null;
        }
    }

    public static class Case2VertexCentricComputeFunction extends AbstractVcFunc<Long, Double, Double, Long> {

        @Override
        public void compute(Long vertexId, Iterator<Long> messageIterator) {
            long iterationId = this.context.getIterationId();

            if (iterationId == 1) {
                List<IEdge<Long, Double>> outEdges = this.context.edges().getOutEdges();
                for (IEdge<Long, Double> outEdge : outEdges) {
                    this.context.sendMessage(outEdge.getTargetId(), vertexId);
                }
            } else if (iterationId == 2) {
                while (messageIterator.hasNext()) {
                    Long msg = messageIterator.next();
                    List<IEdge<Long, Double>> outEdges = this.context.edges().getOutEdges();
                    for (IEdge<Long, Double> outEdge : outEdges) {
                        this.context.sendMessage(outEdge.getTargetId(), msg);
                    }
                }
            } else if (iterationId == 3) {
                Map<Long, Integer> map = new HashMap<>();
                List<IEdge<Long, Double>> outEdges = this.context.edges().getOutEdges();
                for (IEdge<Long, Double> outEdge : outEdges) {
                    Long targetId = outEdge.getTargetId();
                    map.put(targetId, map.getOrDefault(targetId, 0) + 1);
                }
                int cnt = 0;
                while (messageIterator.hasNext()) {
                    Long msg = messageIterator.next();
                    if (map.containsKey(msg)) {
                        cnt += map.get(msg);
                    }
                }
                this.context.setNewVertexValue((double) cnt);
            }
        }
    }
}
